import os
from random import SystemRandom

from util import img_util


# from setting.grass_and_earth import img_dic, img_size, tag_connect_dict, generator_name, asset_img_path
from setting.city import img_dic, img_size, tag_connect_dict, generator_name, asset_img_path

result_size = (10, 10)
out_path = f'out/{generator_name}/'
out_num = 10

rand = SystemRandom()


def init_img():
    '''
    初始化结果集、坍塌位置
    :return:
    '''

    arr = [["" for _ in range(result_size[1])] for _ in range(result_size[0])]
    first_pos = (rand.randint(0, result_size[0] - 1), rand.randint(0, result_size[1] - 1))
    curr_img = rand.choice([i for i in img_dic.keys()])
    # curr_img = "26_11"
    arr[first_pos[0]][first_pos[1]] = set([curr_img])
    part_sure_pos_set = {first_pos}
    return (arr, part_sure_pos_set)

def collapse(res_arr, starting_collapse_pos_set:set, directions_tag_img_dict) -> bool:
    """
    彻底坍塌一个单元格， 并坍塌周围的单元格
    :param res_arr: 结果集
    :param starting_collapse_pos_set: 已经开始坍塌的位置
    :param directions_tag_img_dict:
    :return: 是否发生了冲突
    """
    _tmp_part_sure_pos_list = list(starting_collapse_pos_set)
    _tmp_part_sure_pos_list.sort(key=lambda a: len(res_arr[a[0]][a[1]]))
    _curr_pos = _tmp_part_sure_pos_list[0]
    starting_collapse_pos_set.remove(_curr_pos)
    _curr = res_arr[_curr_pos[0]][_curr_pos[1]]

    if len(_curr) == 0:
        # 没有坍塌的方向， 生成错误
        return False
    if isinstance(_curr, str):
        raise Exception("ERROR!")
    else:
        # 使用类似轮盘赌的方式，选择坍塌方向
        img_choices = list(_curr)
        img_choices_with_scope = []
        _pie_area_total = 0
        for _img_name in img_choices:
            _scope_from = _pie_area_total
            _pie_area_total += img_dic[_img_name]['weight']
            _scope_to = _pie_area_total
            img_choices_with_scope.append((_img_name, (_scope_from, _scope_to)))
        _point = rand.uniform(0, _pie_area_total)
        _curr = ""
        for _img_with_scope in img_choices_with_scope:
            if _point <= _img_with_scope[1][1]:
                _curr = _img_with_scope[0]
        res_arr[_curr_pos[0]][_curr_pos[1]] = _curr

    # 坍塌周围的单元格
    for _direction in range(4):
        if _direction == 0:
            _match_direction = 1
            _next_pos = (_curr_pos[0] - 1, _curr_pos[1])
        elif _direction == 1:
            _match_direction = 0
            _next_pos = (_curr_pos[0] + 1, _curr_pos[1])
        elif _direction == 2:
            _match_direction = 3
            _next_pos = (_curr_pos[0], _curr_pos[1] - 1)
        elif _direction == 3:
            _match_direction = 2
            _next_pos = (_curr_pos[0], _curr_pos[1] + 1)

        # 检查边界
        if _next_pos[0] >= 0 and _next_pos[0] < result_size[0] and _next_pos[1] >= 0 and _next_pos[1] < result_size[1]:
            # 当前单元格在这个方向上的tag
            _curr_tag = img_dic[_curr]['tags'][_direction]
            # 周围单元格的信息
            _next = res_arr[_next_pos[0]][_next_pos[1]]
            # 周围单元在相反方向的tag
            _next_tags = tag_connect_dict[_curr_tag]

            # 根据当前单元格，对周围单元格进行坍塌
            _next_should = set()
            for _tag in _next_tags:
                for _img_name in directions_tag_img_dict[_match_direction][_tag]:
                    _next_should.add(_img_name)

            if _next == "":
                # 还未开始坍塌，初始化坍塌
                res_arr[_next_pos[0]][_next_pos[1]] = _next_should
                starting_collapse_pos_set.add(_next_pos)
            # elif type(_next) == type({1}):
            elif isinstance(_next, set):
                # 已经开始坍塌了，则继续坍塌
                res_arr[_next_pos[0]][_next_pos[1]] = _next.intersection(_next_should)
                starting_collapse_pos_set.add(_next_pos)
    return True



if __name__ == '__main__':
    os.makedirs(name=out_path, exist_ok=True)

    tag_list = tag_connect_dict.keys()
    directions_tag_img_dict = {_direction: {tag: set() for tag in tag_list} for _direction in range(4)}

    for img_name, img_args in img_dic.items():
        for tag_index in range(len(directions_tag_img_dict)):
            img_tag = img_args['tags'][tag_index]
            directions_tag_img_dict[tag_index][img_tag].add(img_name)

    for out_index in range(out_num):
        # result_arr: 保存结果，存储每个单元格使用哪个(些)图片
        # part_sure_pos_set: 开始坍塌的位置集合
        result_arr, starting_collapse_pos_set = init_img()

        # 若没有坍塌的地方，则说明生成完毕
        while len(starting_collapse_pos_set) > 0:
            # 若发生冲突，则重新初始化，重新生成
            if not collapse(result_arr, starting_collapse_pos_set, directions_tag_img_dict):
                print("generate failed, retry...")
                result_arr, starting_collapse_pos_set = init_img()


        # for i in result_arr:
        #     for j in i:
        #         print(j)

        # ########SAVE IMG#########
        out_img_path = f'{out_path}/{out_index}.png'
        img_util.save_img_from_res_arr(img_size, asset_img_path, result_arr, out_img_path, 'png')
        print(f'generated success: {out_img_path}')



